Add BF16 support for int8_vectorwise_quant / LLM.int8 activation quant#1985
Add BF16 support for int8_vectorwise_quant / LLM.int8 activation quant#1985kru2710shna wants to merge 4 commits into
Conversation
| // One block per row. | ||
| // Threads load column values in a striped arrangement. | ||
| // e.g. t0 reads row[0], row[0+nthreads], .. | ||
| // and t1 reads row[1], row[1+nthreads], .. | ||
| // Each thread will determine its local absmax. | ||
| // We then do a blockwise reduction to determine the row's absmax. |
There was a problem hiding this comment.
Why was this comment removed?
There was a problem hiding this comment.
Not intentional. These got dropped when I rewrote the kernel body for float accumulation. They're accurate and unrelated to the change, so I will restored them verbatim in the next commit
| // For sparse decomposition, we do not want to quantize the outliers. | ||
| // Instead they're zeroed out. |
There was a problem hiding this comment.
Same cause as above, accidental during the float-reduction rewrite. Restored. Thanks for catching both.
There was a problem hiding this comment.
This warning would likely no longer apply.
|
I think it would be worthwhile to update some of the other relevant tests to exercise the new bf16 path. E.g. in |
|
Will do |
Templates int8VectorQuant on T and adds bf16 kernel instantiations plus a cint8_vector_quant_bf16 C ABI entry point, mirroring the existing gemm_4bit_inference_naive fp16/bf16/fp32 pattern. The blockwise absmax reduction now accumulates in float rather than T: required for bf16 to compile cleanly and slightly improves fp16 accuracy (rowStats was already float, so downstream is unaffected). Removes the forced A.to(torch.float16) casts in MatMul8bitLt so bf16 activations quantize natively. Closes bitsandbytes-foundation#1868.
…rning - Restore explanatory comments in kInt8VectorQuant that were dropped during the float-accumulation rewrite (striped-load pattern, sparse-decomp absmax, outlier zeroing). - Remove the 'inputs will be cast to float16' warning in MatMul8bitLt, which no longer applies now that the forced fp16 cast is gone.
kru2710shna
force-pushed
the
feature/int8-bf16-vectorwise-quant
branch
from
e30d08c to
cdd9e23
Compare
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
Extends existing tests to exercise the new bf16 activation-quant path: - test_matmullt: drop the forced fp16 cast so the existing bf16 param actually exercises the bf16 kernel end-to-end. - test_int8_linear_matmul_half: parametrize over fp16/bf16. - test_linear8bitlt_inference and new test_linear8bitlt_forward_dtypes: exercise the Linear8bitLt module forward in both dtypes. Addresses review feedback on bitsandbytes-foundation#1985.
|
Done in #1985 - e888a86. Added bf16 coverage across the relevant tests: test_matmullt: dropped the forced fp16 cast so the existing bf16 param now exercises the real path end-to-end. Note: test_linear8bit.py didn't exist, so I covered test_linear8bitlt.py instead. Also rebased onto latest main. Happy to adjust anything. |
|
There's still a test failure and a couple build issues to resolve.
|
Templates kdequant_mm_int32_fp16 and its launcher on output dtype, adds a cdequant_mm_int32_bf16 C ABI entry point, and dispatches on dtype in the Python backend. Fused bias now works for both fp16 and bf16 when the bias dtype matches the output dtype. Exposes a dtype kwarg on bitsandbytes.functional.int8_mm_dequant so callers can request bf16 output directly instead of relying on a cast afterward. Reverts the earlier test workaround now that int8_mm_dequant genuinely returns bf16 when requested. Addresses review feedback on bitsandbytes-foundation#1985 (previously deferred as a follow-up).
|
Went ahead and implemented this properly rather than just patching the test. int8_mm_dequant now natively supports bfloat16 output end-to-end: Templated kdequant_mm_int32_fp16 and its launcher on output dtype, added a bnb_bfloat16 instantiation and a cdequant_mm_int32_bf16 C ABI entry. This should close the last known gap for bf16 support in this PR. |
2 participants
Templates int8VectorQuant on T and adds bf16 kernel instantiations plus a cint8_vector_quant_bf16 C ABI entry point, mirroring the existing gemm_4bit_inference_naive fp16/bf16/fp32 pattern. The blockwise absmax reduction now accumulates in float rather than T: required for bf16 to compile cleanly and slightly improves fp16 accuracy (rowStats was already float, so downstream is unaffected). Removes the forced A.to(torch.float16) casts in MatMul8bitLt so bf16 activations quantize natively.
Closes #1868.